import torch
from model import LSTM
from train import train_model
from test import test_model
from data_processing import load_data
import matplotlib.pyplot as plt

# 配置参数
input_size = 3  # 输入特征数量（根据选择的列数）
hidden_size = 50
num_layers = 2
output_size = 1
num_epochs = 100
learning_rate = 0.001
seq_length = 20
file_path = 'data.csv'
feature_cols = ['feature1', 'feature2', 'feature3']  # 替换为实际列名
target_col = 'target'

# 数据加载
X, y, scaler = load_data(file_path, feature_cols, target_col, seq_length)

# 数据划分
train_size = int(0.8 * len(X))
X_train, y_train = X[:train_size], y[:train_size]
X_test, y_test = X[train_size:], y[train_size:]

# 转换为 PyTorch 张量
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32)

# 实例化模型
model = LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, output_size=output_size)

# 训练模型
train_losses = train_model(model, X_train, y_train, num_epochs, learning_rate)

# 测试模型
y_test_inv, predictions_inv = test_model(model, X_test, y_test, scaler)

# 可视化结果
plt.figure(figsize=(10, 5))
plt.plot(y_test_inv, label='True Values')
plt.plot(predictions_inv, label='Predictions')
plt.legend()
plt.title('True vs Predicted Values')
plt.show()
